#!/bin/bash

PYTHON=python
FILE=CIFAR_train.py

# message
MSG_BATCH=1
MSG_BATCH_FREQ=100

# log
LOG_NAME="CIFAR_imb"
LOG_EPOCH_FREQ=1
LOG_BATCH_FREQ=1000

# data
BATCH=8
WORKERS=1
EPOCH=35

# model
MODEL="resnet34"
SIZE=224
CLASSES=50030
LOAD=""
TORCH_PRETRAIN=1
SAVE_DIR="/home/ruideng/workspace/LSPR/CIFAR_weights"
SAVE_FREQ=10

# eval
EVAL_BATCH_FREQ=0
EVAL_EPOCH_FREQ=5

# scheduler
SCHEDULER=""
MILESTONES=""
GAMMA=1e-1

# optimizer
OPTIM=adam
LR=1e-2
BETAS="0.95 0.999"
WEIGHT_DECAY=0.01
MOMENTUM=0.9

# resume
RESUME=0

# GPU
DEVICE=cuda

ARG=' '

# message
if [ ${MSG_BATCH} -eq 1 ]
then
    ARG="${ARG} --msg_batch"
    ARG="${ARG} --msg_batch_freq ${MSG_BATCH_FREQ}"
fi

# log
if [ ${LOG_NAME} ]
then
    ARG="${ARG} --log_name ${LOG_NAME}"
fi
if [ ${LOG_EPOCH_FREQ} -ne 0 ]
then
    ARG="${ARG} --log_epoch_freq ${LOG_EPOCH_FREQ}"
fi
if [ ${LOG_BATCH_FREQ} -ne 0 ]
then
    ARG="${ARG} --log_batch_freq ${LOG_BATCH_FREQ}"
fi

# data
ARG="${ARG} --batch ${BATCH}"
ARG="${ARG} --workers ${WORKERS}"
ARG="${ARG} --epoch ${EPOCH}"

# model
if [ ${MODEL} ]
then
    ARG="${ARG} --model ${MODEL}"
fi
ARG="${ARG} --size ${SIZE}"
ARG="${ARG} --classes ${CLASSES}"
if [ ${LOAD} ]
then
    ARG="${ARG} --load ${LOAD}"
fi
if [ ${TORCH_PRETRAIN} -eq 1 ]
then
    ARG="${ARG} --torch_pretrain"
fi
if [ ${SAVE_DIR} ]
then
    ARG="${ARG} --save_dir ${SAVE_DIR}"
    ARG="${ARG} --save_freq ${SAVE_FREQ}"
fi

# eval
ARG="${ARG} --eval_batch_freq ${EVAL_BATCH_FREQ}"
ARG="${ARG} --eval_epoch_freq ${EVAL_EPOCH_FREQ}"

# scheduler
if [ ${SCHEDULER} ]
then
    ARG="${ARG} --scheduler ${SCHEDULER}"
    ARG="${ARG} --milestones ${MILESTONES}"
    ARG="${ARG} --gamma ${GAMMA}"
fi

# optimizer
ARG="${ARG} --lr ${LR}"
ARG="${ARG} --optim ${OPTIM}"
ARG="${ARG} --betas ${BETAS}"
ARG="${ARG} --weight_decay ${WEIGHT_DECAY}"
ARG="${ARG} --momentum ${MOMENTUM}"

# resume
if [ ${RESUME} -eq 1 ]
then
    ARG="${ARG} --resume"
fi

# GPU
ARG="${ARG} --device ${DEVICE}"

# echo ${ARG}

echo "$PYTHON $FILE $ARG"

# $PYTHON $FILE $ARG
